from fastapi import APIRouter, HTTPException, Depends, Query
from typing import List

from database.database import Document, VectorEmbedding, KnowledgeBase
from database.database_function import get_db
from api.schemas.knowledge_base import KnowledgeBaseCreateRequest, KnowledgeBaseUpdateRequest, KnowledgeBaseResponse, KnowledgeBaseListResponse


router = APIRouter()


@router.post("/knowledge-bases", response_model=KnowledgeBaseResponse, summary="创建知识库")
async def create_knowledge_base(body: KnowledgeBaseCreateRequest, db=Depends(get_db)):
    existing = db.query(KnowledgeBase).filter(KnowledgeBase.name == body.name).first()
    if existing:
        raise HTTPException(status_code=400, detail="知识库名称已存在")
    kb = KnowledgeBase(name=body.name, description=body.description)
    if getattr(body, 'embedding_model_name', None):
        kb.embedding_model_name = body.embedding_model_name
    db.add(kb)
    db.commit()
    db.refresh(kb)
    return KnowledgeBaseResponse(
        id=kb.id,
        name=kb.name,
        description=kb.description,
        embedding_model_name=getattr(kb, 'embedding_model_name', None),
        created_at=kb.created_at,
        updated_at=kb.updated_at,
        document_count=0
    )


@router.get("/knowledge-bases", response_model=KnowledgeBaseListResponse, summary="知识库列表")
async def list_knowledge_bases(
    page: int = Query(1, ge=1, description="页码"),
    page_size: int = Query(10, ge=1, le=100, description="每页数量"),
    db=Depends(get_db)
):
    offset = (page - 1) * page_size
    q = db.query(KnowledgeBase)
    total = q.count()
    items = q.offset(offset).limit(page_size).all()
    responses: List[KnowledgeBaseResponse] = []
    for kb in items:
        doc_count = db.query(Document).filter(Document.kb_id == kb.id).count()
        responses.append(KnowledgeBaseResponse(
            id=kb.id,
            name=kb.name,
            description=kb.description,
            embedding_model_name=getattr(kb, 'embedding_model_name', None),
            created_at=kb.created_at,
            updated_at=kb.updated_at,
            document_count=doc_count
        ))
    return KnowledgeBaseListResponse(
        success=True,
        knowledge_bases=responses,
        total_count=total,
        page=page,
        page_size=page_size,
        total_pages=(total + page_size - 1) // page_size
    )


@router.get("/knowledge-bases/{kb_id}", response_model=KnowledgeBaseResponse, summary="知识库详情")
async def get_knowledge_base(kb_id: int, db=Depends(get_db)):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    doc_count = db.query(Document).filter(Document.kb_id == kb.id).count()
    return KnowledgeBaseResponse(
        id=kb.id,
        name=kb.name,
        description=kb.description,
        embedding_model_name=getattr(kb, 'embedding_model_name', None),
        created_at=kb.created_at,
        updated_at=kb.updated_at,
        document_count=doc_count
    )


@router.patch("/knowledge-bases/{kb_id}", response_model=KnowledgeBaseResponse, summary="更新知识库")
async def update_knowledge_base(kb_id: int, body: KnowledgeBaseUpdateRequest, db=Depends(get_db)):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    if body.name:
        exists = db.query(KnowledgeBase).filter(KnowledgeBase.name == body.name, KnowledgeBase.id != kb_id).first()
        if exists:
            raise HTTPException(status_code=400, detail="知识库名称已存在")
        kb.name = body.name
    if body.description is not None:
        kb.description = body.description
    if getattr(body, 'embedding_model_name', None) is not None:
        kb.embedding_model_name = body.embedding_model_name
    db.commit()
    db.refresh(kb)
    doc_count = db.query(Document).filter(Document.kb_id == kb.id).count()
    return KnowledgeBaseResponse(
        id=kb.id,
        name=kb.name,
        description=kb.description,
        embedding_model_name=getattr(kb, 'embedding_model_name', None),
        created_at=kb.created_at,
        updated_at=kb.updated_at,
        document_count=doc_count
    )


@router.delete("/knowledge-bases/{kb_id}", summary="删除知识库")
async def delete_knowledge_base(kb_id: int, db=Depends(get_db)):
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    doc_ids = [d.id for d in db.query(Document.id).filter(Document.kb_id == kb_id).all()]
    if doc_ids:
        db.query(VectorEmbedding).filter(VectorEmbedding.doc_id.in_(doc_ids)).delete(synchronize_session=False)
        db.query(Document).filter(Document.id.in_(doc_ids)).delete(synchronize_session=False)
    db.delete(kb)
    db.commit()
    return {"success": True, "message": "知识库及其文档已删除"}


